import os
import time
import torch
from diffusers import StableDiffusionPipeline
from peft import PeftModel
from safetensors.torch import load_file

def log(message):
    """打印带时间戳的日志"""
    from datetime import datetime
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f"[{timestamp}] {message}")

def main():
    start_time = time.time()
    log("开始处理...")

    try:
        # 设置模型路径
        base_model_path = 'E:/workspace/llm/text2image/models/AI-ModelScope/stable-diffusion-v1-5'
        lora_model_path = 'E:/workspace/llm/text2image/models/tangmuahu/Lora_Test'

        # 检查模型文件
        log(f"检查模型路径...")
        if not os.path.exists(base_model_path):
            raise FileNotFoundError(f"基础模型路径不存在: {base_model_path}")
        if not os.path.exists(lora_model_path):
            raise FileNotFoundError(f"LoRA模型路径不存在: {lora_model_path}")

        # 设置设备
        device = "cuda" if torch.cuda.is_available() else "cpu"
        log(f"使用设备: {device}")
        if device == "cuda":
            log(f"CUDA设备: {torch.cuda.get_device_name(0)}")
            log(f"CUDA可用内存: {torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB")

        # 加载基础模型
        log("加载基础模型...")
        pipe = StableDiffusionPipeline.from_pretrained(
            base_model_path,
            torch_dtype=torch.float32,  # 使用float32以提高兼容性
            safety_checker=None,
            requires_safety_checker=False
        ).to(device)

        # 加载LoRA权重
        log("加载LoRA权重...")
        lora_path = os.path.join(lora_model_path, 'epoch=0-step=500.ckpt')
        if not os.path.exists(lora_path):
            raise FileNotFoundError(f"LoRA文件不存在: {lora_path}")

        # 加载checkpoint
        state_dict = torch.load(lora_path, map_location='cpu')
        if "state_dict" in state_dict:
            state_dict = state_dict["state_dict"]

        # 应用LoRA权重到UNet
        alpha = 0.75
        for key in state_dict:
            if "lora_unet" in key:
                base_key = key.replace("lora_unet.", "").replace("_lora", "")
                if "lora_down" in key:
                    down_weight = state_dict[key]
                    up_key = key.replace("lora_down", "lora_up")
                    up_weight = state_dict[up_key]

                    # 找到对应的原始层
                    current_layer = pipe.unet
                    for part in base_key.split("."):
                        if part.isdigit():
                            part = int(part)
                        current_layer = getattr(current_layer, part)

                    if hasattr(current_layer, "weight"):
                        # 计算并应用LoRA更新
                        weight = current_layer.weight
                        if len(weight.shape) == 2:
                            # 线性层
                            delta = (up_weight @ down_weight) * alpha
                        else:
                            # 卷积层
                            delta = torch.nn.functional.conv2d(
                                down_weight.permute(1, 0, 2, 3),
                                up_weight.permute(1, 0, 2, 3)
                            ).permute(1, 0, 2, 3) * alpha

                        weight.data += delta.to(weight.device)

        # 生成图像
        log("生成图像中...")
        prompt = "A beautiful girl in the sunshine, high quality, detailed"
        negative_prompt = "low quality, blurry, worst quality, deformed"

        # 生成图像
        with torch.inference_mode():
            result = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=30,
                guidance_scale=7.5,
            )

        # 保存图像
        output_path = "result11.png"
        result.images[0].save(output_path)
        log(f"图像已保存到: {output_path}")

        # 显示图像信息
        width, height = result.images[0].size
        log(f"图像尺寸: {width}x{height}")
        log(f"图像模式: {result.images[0].mode}")

        # 检查图像是否为全黑
        image_array = torch.tensor(list(result.images[0].getdata())).float()
        avg_pixel = image_array.mean().item()
        log(f"平均像素值: {avg_pixel:.2f} (0=黑色, 255=白色)")

        # 清理资源
        del pipe
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        # 计算总时间
        elapsed_time = time.time() - start_time
        log(f"处理完成! 总耗时: {elapsed_time:.2f}秒")

    except Exception as e:
        log(f"发生错误: {str(e)}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()